package com.cl.ks.flow.handler;

import com.alibaba.fastjson.JSONObject;
import com.cl.ks.config.ValueConfig;
import com.cl.ks.entity.KsDataField;
import com.cl.ks.flow.base.BaseFlowNodeHandler;
import com.cl.ks.flow.enums.NodeHandlerEnum;
import com.cl.ks.utils.DbUtil;
import com.cl.spark.dto.SparkParam;
import com.cl.spark.dto.SparkResult;
import com.jandar.ds.pojo.DbInfo;
import com.querydsl.jpa.impl.JPAQueryFactory;
import org.apache.spark.rdd.RDD;
import org.apache.spark.sql.*;
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.Metadata;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.util.*;
import java.util.stream.Collectors;

import static com.cl.ks.constant.QTableConstant.*;


@Component
public class DataTypeHandler extends BaseFlowNodeHandler {
    private final DbUtil dbUtil;
    private final JPAQueryFactory queryFactory;

    private final SparkSession sparkSession;

    private final ValueConfig valueConfig;


    public DataTypeHandler(DbUtil dbUtil, JPAQueryFactory queryFactory, SparkSession sparkSession, ValueConfig valueConfig) {
        this.dbUtil = dbUtil;
        this.queryFactory = queryFactory;
        this.sparkSession = sparkSession;
        this.valueConfig = valueConfig;
    }

    @Override
    public SparkResult process(ProcessParam processParam) {


        JSONObject expression = processParam.getSparkNode().getNodeExpression();
        String sparkNodeName = "";
        String tableName = String.valueOf(expression.get("targetTable"));
        String dbCode = String.valueOf(expression.get("dbCode"));
        String ksDataTypeCode = String.valueOf(expression.get("code"));
        sparkNodeName = "数据资源：" + dbCode + "-" + tableName;

        if (tableName == null || dbCode == null) {
            throw new RuntimeException(ksDataTypeCode + "未配置库表");
        }

        List<String> fieldsAs = new ArrayList<>();

        List<String> fields = new ArrayList<>();
        List<String> fieldsAlias = new ArrayList<>();
        List<Map<String, Object>> columns = new ArrayList<>();
        if (!dbCode.equals("null") && !tableName.equals("null")) {
            String profilesActive = valueConfig.getProfilesActive();
            if ("dm".equals(profilesActive)) {
                columns = dbUtil.fetchSQL(dbCode, "select COLUMN_NAME as Field  from  USER_TAB_COLUMNS where TABLE_NAME='" + tableName + "'");
            } else {
                columns = dbUtil.fetchSQL(dbCode, "show full columns from " + tableName);
            }
        }
        List<String> fieldList = columns.stream().map(item -> String.valueOf(item.get("Field"))).collect(Collectors.toList());

        if (ksDataTypeCode != null) {
            processParam.getKsFlowGraph().getKsDataTypeCodeList().add(ksDataTypeCode);
            List<Integer> sysSourceConfigIds = queryFactory.select(qSysSourceConfig.id).from(qSysSourceConfig).where(qSysSourceConfig.dbCode.eq(dbCode).and(qSysSourceConfig.targetTable.eq(tableName))).fetch();
            processParam.getKsFlowGraph().getSourceConfigIdList().addAll(sysSourceConfigIds.stream().map(String::valueOf).collect(Collectors.toList()));


            List<KsDataField> ksDataFields = queryFactory.selectFrom(qKsDataField).where(qKsDataField.dataTypeCode.eq(ksDataTypeCode)).fetch();
            ksDataFields.stream()
                    .filter(item -> {

                                if (fieldList.isEmpty() || "ruleName".equals(item.getFieldMapping())) {
                                    return true;
                                } else {
                                    return fieldList.stream().anyMatch(i -> i.equalsIgnoreCase(item.getCode()));
                                }
                            }
                    )
                    .forEach(ksDataField -> {
                        String tableField = ksDataField.getCode();
                        String asName = ksDataField.getFieldMapping();
                        if (asName != null) {
                            Optional<String> exist = fieldsAlias.stream().filter(asName::equalsIgnoreCase).findFirst();
                            if (exist.isPresent()) {
                                asName = asName + "_repeat";
                            }
                        } else {
                            asName = tableField;
                        }
                        fields.add(tableField);
                        fieldsAlias.add(asName.trim());
                    });
        }

        fieldList.stream().filter(i -> fieldsAlias.stream().noneMatch(i::equalsIgnoreCase)).forEach(i -> {
            fields.add(i);
            fieldsAlias.add(i.trim());
        });


        Optional<SparkResult> originalSparkResultOptional = processParam.getSparkNode().getInputSparkResultList().values().stream().findFirst();

        if (originalSparkResultOptional.isPresent()) {
            SparkResult sparkResult = originalSparkResultOptional.get();
            Dataset<Row> dataset = sparkResult.getDataset();
            if (dataset != null && !dataset.isEmpty()) {
                for (int i = 0; i < dataset.columns().length; i++) {
                    int fieldsIndex = fieldsAlias.indexOf(dataset.columns()[i]);
                    if (fieldsIndex != -1) {
                        fieldsAs.add(fieldsAlias.get(fieldsIndex) + " as " + fields.get(fieldsIndex));
                    }
                }
                Dataset<Row> newDataset = dataset.selectExpr(fieldsAs.toArray(new String[0]));
                sparkResult.setDataset(newDataset);
                return sparkResult;
            }

        }

        for (int i = 0; i < fields.size(); i++) {
            fieldsAs.add(fields.get(i) + " as " + fieldsAlias.get(i));
        }

        Integer lastPosition = 0;
        if (processParam.getKsFlowGraph().getIncrementFlag() != null && processParam.getKsFlowGraph().getIncrementFlag()) {
            lastPosition = queryFactory.select(qKsFlowGraphLog.lastPosition).from(qKsFlowGraphLog).where(qKsFlowGraphLog.ksFlowGraphId.eq(processParam.getKsFlowGraph().getId()).and(qKsFlowGraphLog.sparkNodeCode.eq(processParam.getSparkNode().getCode()))).orderBy(qKsFlowGraphLog.ksFlowGraphRecordId.desc()).fetchFirst();
            if (lastPosition == null) {
                lastPosition = 0;
            }
        }

        if (processParam.getTestFlag() == true || StringUtils.hasLength(processParam.getKsFlowGraph().getFailedNodeCode())) {
            StructType schema = new StructType(
                    fieldsAlias.stream()
                            .filter(fieldName -> fieldName != null && !fieldName.trim().isEmpty())
                            .map(fieldName -> new StructField(fieldName, DataTypes.StringType, true, Metadata.empty()))
                            .toArray(StructField[]::new)
            );
            Dataset<Row> emptyDataFrame = sparkSession.createDataFrame(new ArrayList<>(), schema);

            return SparkResult.success(emptyDataFrame);
        }

        Long count = Long.parseLong(String.valueOf(dbUtil.fetchSQL(dbCode, "select count(1) as count from " + tableName).get(0).get("count")));
        long limit = count - lastPosition;
        processParam.getKsFlowGraph().setSourceConfigDataTotal((int) (processParam.getKsFlowGraph().getSourceConfigDataTotal() + (count - lastPosition)));
        processParam.getKsFlowGraphLog().setLastPosition(Math.toIntExact(count));

        DbInfo dbInfo = dbUtil.getDbInfo(dbCode);
        SparkParam sparkParam = new SparkParam();
        expression.put("limit", limit);
        expression.put("skip", lastPosition);
        expression.put("url", dbInfo.getUrl());
        expression.put("username", dbInfo.getUsername());
        expression.put("password", dbInfo.getPassword());
        expression.put("tableName", tableName);
        expression.put("fieldsAs", String.join(",", fieldsAs));
        sparkParam.setNodeExpression(expression);
        return sparkNodeFactory.getSparkNodeByCode("dataSource").handle(sparkParam);
    }


    @Override
    public NodeHandlerEnum getType() {
        return NodeHandlerEnum.DATA_TYPE;
    }
}
